from torch import nn
import torch
from chinese_text_classification.utils import get_all_need
from chinese_text_classification.model import Model
from chinese_text_classification.train_dve import train

train_loader, dev_loader, embeddings_matrix = get_all_need(35, 64, 128)  # 得到模型需要的数据
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
my_model = Model(embeddings_matrix, 128, [3, 4, 5], 0.5, 10).to(device)  # 实例化模型
print(my_model)
my_optimizer = torch.optim.Adam(my_model.parameters(), lr=0.001)  # 优化器
my_loss_func = nn.CrossEntropyLoss()  # 损失函数
train(train_loader, dev_loader, my_model, my_loss_func, my_optimizer, 30)  # 训练
